{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Working with multiple objects in the Python SDK\n",
    "AIS supports multi-object operations on groups of objects. An `ObjectGroup` can be created with one of:\n",
    "- a list of object names\n",
    "- an [ObjectRange](https://github.com/NVIDIA/aistore/blob/main/python/aistore/sdk/multiobj/object_range.py)\n",
    "- a string template."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install aistore"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Set up the client and create necessary buckets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from aistore import Client\n",
    "from aistore.sdk.errors import AISError\n",
    "import os\n",
    "\n",
    "ais_url = os.getenv(\"AIS_ENDPOINT\", \"http://localhost:8080\")\n",
    "client = Client(ais_url)\n",
    "bucket = client.bucket(\"my-bck\").create(exist_ok=True)\n",
    "copy_dest_bucket = client.bucket(\"copy-destination-bucket\").create(exist_ok=True)\n",
    "transform_dest_bucket = client.bucket(\"transform-destination-bucket\").create(\n",
    "    exist_ok=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Create some objects in the bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "object_names = [f\"example_obj_{i}\" for i in range(10)]\n",
    "for name in object_names:\n",
    "    bucket.object(name).get_writer().put_content(\"object content\".encode(\"utf-8\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Create Object Group by list of names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "my_objects = bucket.objects(obj_names=object_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Create Object Group by ObjectRange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from aistore.sdk.multiobj import ObjectRange\n",
    "\n",
    "my_object_range = ObjectRange(prefix=\"example_obj_\", min_index=1, max_index=3)\n",
    "my_objects = bucket.objects(obj_range=my_object_range)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Create Object Group by Template String\n",
    "String templates can be passed directly to AIS following the [syntax described here](https://github.com/NVIDIA/aistore/blob/main/docs/batch.md#operations-on-multiple-selected-objects)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Equivalent to the range above\n",
    "my_object_template = \"example_obj_{1..3}\"\n",
    "my_objects = bucket.objects(obj_template=my_object_template)\n",
    "# More advanced template example with multiple ranges and defined steps\n",
    "complex_range = \"example_obj_{0..10..2}_details_{1..9..2}.file-extension\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Prefetch or evict multiple objects when using a bucket with a cloud backend"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "my_objects.prefetch()\n",
    "my_objects.evict()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Copy multiple objects\n",
    "\n",
    "Copies selected objects directly to the new bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "copy_job = my_objects.copy(to_bck=copy_dest_bucket)\n",
    "# The job will reach an idle state before finishing, so wait for idle\n",
    "client.job(job_id=copy_job).wait_for_idle()\n",
    "# See the objects in the destination bucket\n",
    "copy_dest_bucket.list_all_objects()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Delete multiple objects from the destination bucket above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "all_objects = copy_dest_bucket.list_all_objects()\n",
    "# Creates a group including all objects from the destination bucket\n",
    "objects_to_delete = copy_dest_bucket.objects(\n",
    "    obj_names=[entry.name for entry in all_objects]\n",
    ")\n",
    "delete_job_id = objects_to_delete.delete()\n",
    "client.job(delete_job_id).wait()\n",
    "after_deletion = copy_dest_bucket.list_all_objects()\n",
    "print(\n",
    "    f\"Objects before deletion: {len(all_objects)}, objects after deletion: {len(after_deletion)}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Transform -- Provide an ETL to be performed on each object so the result appears in the destination bucket.\n",
    "\n",
    "Note: This step requires the AIS cluster to be running in Kubernetes; see [getting_started](https://github.com/NVIDIA/aistore/blob/main/docs/getting_started.md#kubernetes-playground) for setup info."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# First create an ETL\n",
    "# This is a simple example transform that reverses each object's contents (assuming utf-8 encoded text)\n",
    "from aistore.sdk.etl.webserver.fastapi_server import FastAPIServer\n",
    "\n",
    "class ReverseTransformETL(FastAPIServer):\n",
    "    def transform(self, data: bytes, path: str, etl_args: str) -> bytes:\n",
    "        \"\"\"Reverse the contents of each object (assuming utf-8 encoded text)\"\"\"\n",
    "        reversed_in_str = data.decode(\"utf-8\")[::-1]\n",
    "        return reversed_in_str.encode()\n",
    "\n",
    "etl_name = \"multiobj-transform-example\"\n",
    "try:\n",
    "    client.etl(etl_name=etl_name).init_class()(ReverseTransformETL)\n",
    "except AISError as err:\n",
    "    print(err)\n",
    "\n",
    "# Now run the transform with the etl name specified\n",
    "transform_job = my_objects.transform(etl_name=etl_name, to_bck=transform_dest_bucket)\n",
    "client.job(job_id=transform_job).wait_for_idle()\n",
    "\n",
    "# The output will be in the destination bucket\n",
    "transformed_objs = transform_dest_bucket.list_all_objects()\n",
    "\n",
    "# See the result\n",
    "for entry in transformed_objs:\n",
    "    input_data = bucket.object(entry.name).get_reader().read_all()\n",
    "    output_data = transform_dest_bucket.object(entry.name).get_reader().read_all()\n",
    "    print(f\"Object {entry.name} {input_data} => {output_data}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Cleanup buckets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "for bck in [bucket, copy_dest_bucket, transform_dest_bucket]:\n",
    "    bck.delete(missing_ok=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.1 (main, Dec  7 2022, 01:11:34) [GCC 11.3.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "ead1b95f633dc9c51826328e1846203f51a198c6fb5f2884a80417ba131d4e82"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
